from langchain.tools import tool
import requests
from langchain.chains import create_sql_query_chain
from langchain_community.chat_models import ChatOpenAI  # 或者使用其他支持的LLM
from langchain_community.utilities import SQLDatabase
from langchain.tools import tool



# 初始化数据库连接（假设库名为xeasy，根据你的实际情况修改连接参数）
db = SQLDatabase.from_uri(
    f"mysql+pymysql://root:root@localhost:3306/db_xeasy",  # 替换为你的连接信息
    sample_rows_in_table_info=2  # 在表信息中示例的行数
)


# 初始化LLM模型
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

# 创建SQL查询链
chain = create_sql_query_chain(llm, db)

# 定义工具
@tool
def query_student_natural_language(question: str) -> str:
    """根据自然语言问题查询student表。例如：‘年龄最大的学生是谁？’、‘女生有多少人？’"""
    try:
        # 让LLM生成SQL并执行
        sql_query = chain.invoke({"question": question})
        print(f"生成的SQL: {sql_query}")  # 调试用
        result = db.run(sql_query)
        return f"问题: '{question}' 的查询结果：{result}"
    except Exception as e:
        return f"查询失败：{str(e)}"